from flask_sqlalchemy import SQLAlchemy
from datetime import datetime
import bcrypt

db = SQLAlchemy()

class User(db.Model):
    """用户模型"""
    __tablename__ = 'users'
    
    id = db.Column(db.Integer, primary_key=True)
    username = db.Column(db.String(80), unique=True, nullable=False)
    password_hash = db.Column(db.String(200), nullable=False)
    is_admin = db.Column(db.Boolean, default=False)
    created_at = db.Column(db.DateTime, default=datetime.utcnow)
    
    def set_password(self, password):
        """设置密码"""
        self.password_hash = bcrypt.hashpw(password.encode('utf-8'), bcrypt.gensalt()).decode('utf-8')
    
    def check_password(self, password):
        """验证密码"""
        return bcrypt.checkpw(password.encode('utf-8'), self.password_hash.encode('utf-8'))
    
    def to_dict(self):
        """转换为字典"""
        return {
            'id': self.id,
            'username': self.username,
            'is_admin': self.is_admin,
            'created_at': self.created_at.isoformat()
        }


class ExtractionRecord(db.Model):
    """知识抽取记录模型"""
    __tablename__ = 'extraction_records'
    
    id = db.Column(db.Integer, primary_key=True)
    user_id = db.Column(db.Integer, db.ForeignKey('users.id'), nullable=False)
    prompt = db.Column(db.Text, nullable=False)
    document = db.Column(db.Text, nullable=False)
    result = db.Column(db.Text, nullable=False)  # JSON格式存储
    created_at = db.Column(db.DateTime, default=datetime.utcnow)
    
    user = db.relationship('User', backref=db.backref('extractions', lazy=True))
    
    def to_dict(self):
        """转换为字典"""
        return {
            'id': self.id,
            'user_id': self.user_id,
            'username': self.user.username,
            'prompt': self.prompt,
            'document': self.document,
            'result': self.result,
            'created_at': self.created_at.isoformat()
        }


class KnowledgeEntity(db.Model):
    """知识实体模型"""
    __tablename__ = 'knowledge_entities'
    
    id = db.Column(db.Integer, primary_key=True)
    extraction_id = db.Column(db.Integer, db.ForeignKey('extraction_records.id'), nullable=False)
    entity_type = db.Column(db.String(50), nullable=False)  # 设备名、性能表征、故障状态、故障原因、解决方法
    entity_value = db.Column(db.String(200), nullable=False)
    created_at = db.Column(db.DateTime, default=datetime.utcnow)
    
    extraction = db.relationship('ExtractionRecord', backref=db.backref('entities', lazy=True))
    
    def to_dict(self):
        """转换为字典"""
        return {
            'id': self.id,
            'extraction_id': self.extraction_id,
            'entity_type': self.entity_type,
            'entity_value': self.entity_value,
            'created_at': self.created_at.isoformat()
        }


class KnowledgeRelation(db.Model):
    """知识关系模型"""
    __tablename__ = 'knowledge_relations'
    
    id = db.Column(db.Integer, primary_key=True)
    extraction_id = db.Column(db.Integer, db.ForeignKey('extraction_records.id'), nullable=False)
    subject = db.Column(db.String(200), nullable=False)
    predicate = db.Column(db.String(100), nullable=False)
    object = db.Column(db.String(200), nullable=False)
    created_at = db.Column(db.DateTime, default=datetime.utcnow)
    
    extraction = db.relationship('ExtractionRecord', backref=db.backref('relations', lazy=True))
    
    def to_dict(self):
        """转换为字典"""
        return {
            'id': self.id,
            'extraction_id': self.extraction_id,
            'subject': self.subject,
            'predicate': self.predicate,
            'object': self.object,
            'created_at': self.created_at.isoformat()
        }


class Document(db.Model):
    """文档管理模型"""
    __tablename__ = 'documents'
    
    id = db.Column(db.Integer, primary_key=True)
    title = db.Column(db.String(200), nullable=False)  # 文档标题
    filename = db.Column(db.String(200), nullable=False)  # 原始文件名
    file_type = db.Column(db.String(50), nullable=False)  # 文件类型（docx, pdf, md等）
    file_path = db.Column(db.String(500), nullable=False)  # 文件存储路径
    content = db.Column(db.Text)  # 解析后的文本内容
    file_size = db.Column(db.Integer)  # 文件大小（字节）
    embedding = db.Column(db.Text)  # 向量嵌入（JSON格式）
    upload_user_id = db.Column(db.Integer, db.ForeignKey('users.id'), nullable=False)  # 上传用户
    created_at = db.Column(db.DateTime, default=datetime.utcnow)
    updated_at = db.Column(db.DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
    
    upload_user = db.relationship('User', backref=db.backref('documents', lazy=True))
    
    def to_dict(self):
        """转换为字典"""
        return {
            'id': self.id,
            'title': self.title,
            'filename': self.filename,
            'file_type': self.file_type,
            'file_size': self.file_size,
            'upload_user': self.upload_user.username if self.upload_user else '未知',
            'upload_user_id': self.upload_user_id,
            'created_at': self.created_at.isoformat(),
            'updated_at': self.updated_at.isoformat()
        }


class DictionaryTerm(db.Model):
    """领域词典模型"""
    __tablename__ = 'dictionary_terms'
    
    id = db.Column(db.Integer, primary_key=True)
    term = db.Column(db.String(200), nullable=False)  # 词条
    category = db.Column(db.String(100), nullable=False)  # 类别
    definition = db.Column(db.Text, nullable=False)  # 定义
    examples = db.Column(db.Text)  # 示例（可选）
    embedding = db.Column(db.Text)  # 向量嵌入（JSON格式）
    created_at = db.Column(db.DateTime, default=datetime.utcnow)
    updated_at = db.Column(db.DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
    
    def to_dict(self):
        """转换为字典"""
        return {
            'id': self.id,
            'term': self.term,
            'category': self.category,
            'definition': self.definition,
            'created_at': self.created_at.isoformat(),
            'updated_at': self.updated_at.isoformat()
        }
    
    def to_knowledge_text(self):
        """转换为知识库文本格式"""
        return f"【{self.category}】{self.term}：{self.definition}"


class Conversation(db.Model):
    """对话会话模型"""
    __tablename__ = 'conversations'
    
    id = db.Column(db.Integer, primary_key=True)
    user_id = db.Column(db.Integer, db.ForeignKey('users.id'), nullable=False)
    title = db.Column(db.String(200), default='新对话')
    created_at = db.Column(db.DateTime, default=datetime.utcnow)
    updated_at = db.Column(db.DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
    
    user = db.relationship('User', backref=db.backref('conversations', lazy=True))
    
    def to_dict(self):
        """转换为字典"""
        return {
            'id': self.id,
            'user_id': self.user_id,
            'username': self.user.username if self.user else '未知用户',
            'title': self.title,
            'created_at': self.created_at.isoformat(),
            'updated_at': self.updated_at.isoformat()
        }


class Message(db.Model):
    """对话消息模型"""
    __tablename__ = 'messages'
    
    id = db.Column(db.Integer, primary_key=True)
    conversation_id = db.Column(db.Integer, db.ForeignKey('conversations.id'), nullable=False)
    role = db.Column(db.String(20), nullable=False)  # user 或 assistant
    content = db.Column(db.Text, nullable=False)
    context = db.Column(db.Text)  # 检索到的上下文（可选）
    created_at = db.Column(db.DateTime, default=datetime.utcnow)
    
    conversation = db.relationship('Conversation', backref=db.backref('messages', lazy=True, order_by='Message.created_at'))
    
    def to_dict(self):
        """转换为字典"""
        return {
            'id': self.id,
            'conversation_id': self.conversation_id,
            'role': self.role,
            'content': self.content,
            'context': self.context,
            'created_at': self.created_at.isoformat()
        }
